package com.devops.admin.controller;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.devops.admin.model.IpAddress;
import com.devops.admin.po.DOSBaseUserPo;
import com.devops.admin.service.IpAddressService;
import com.devops.shiro.service.DOSBaseUserService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.kafka.annotation.KafkaListener;
import org.springframework.kafka.support.Acknowledgment;
import org.springframework.kafka.support.KafkaHeaders;
import org.springframework.messaging.handler.annotation.Header;
import org.springframework.messaging.handler.annotation.Payload;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.stream.Collectors;

@ServerEndpoint("/websocket/request/logs/malicious")
@Component
@Slf4j
public class RequestLogsWSController {

    private static IpAddressService ipAddressService;
    private static DOSBaseUserService authService;

    @Autowired
    public void setIpAddressService(IpAddressService service) {
        RequestLogsWSController.ipAddressService = service;
    }

    @Autowired
    public void setAuthService(DOSBaseUserService service) {
        RequestLogsWSController.authService = service;
    }

    @Value(value = "${jysd.office.ips:117.158.0.19,171.15.162.238}")
    private String[] officeIps;

    private static final CopyOnWriteArraySet<Session> sessions = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session) {
        log.info("WebSocket connection opened: " + session.getId());
    }

    @OnClose
    public void onClose(Session session) {
        sessions.remove(session);
        log.info("WebSocket connection closed: " + session.getId());
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("WebSocket error: " + error.getMessage());
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        if (JSONUtil.isTypeJSON(message)) {
            JSONObject json = JSONUtil.parseObj(message);
            if (json.getStr("type").equals("authentication")) {
                String token = json.getStr("token");
                List<DOSBaseUserPo> userPoList = authService.getUserByApiToken(token);
                if (CollUtil.isNotEmpty(userPoList)) {
                    sessions.add(session);
                    sendMessage(session, String.format("Welcome %s!", userPoList.get(0).getLoginName()));
                    log.info("New authenticated WebSocket connection: " + session.getId());
                }
            }
        }
    }

    private void sendMessage(Session session, String message) {
        try {
            session.getBasicRemote().sendText(message);
        } catch (IOException e) {
            log.error("Error sending message: " + e.getMessage());
        }
    }

    public static void broadcastMessage(String message) {
        for (Session session : sessions) {
            if (session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (IOException e) {
                    log.error("Error broadcasting message: " + e.getMessage());
                }
            }
        }
    }

    @KafkaListener(topics = "${kafka.topic.request.logs:public.apisix.request.logs}", concurrency = "3")
    public void consumeTerminalUseRecord(@Payload String message, @Header(KafkaHeaders.ACKNOWLEDGMENT) Acknowledgment ack) {
        try {
            JSONArray filtered = filter(message);
            filtered.forEach(item -> broadcastMessage(JSONUtil.toJsonPrettyStr(item.toString())));
        } catch (Exception e) {
            log.error("Error processing message: " + e.getMessage());
        }
        ack.acknowledge();
    }

    private JSONArray filter(String message) {
        JSONArray inputArray = JSONUtil.parseArray(message);
        return new JSONArray(
                inputArray.stream()
                        .map(JSONObject.class::cast)
                        .filter(item -> item.getByPath("response.status", Integer.class) == 404)
                        //.filter(item -> !isExcludedIp(item.getStr("client_ip")))
                        .map(this::processItem)
                        //.filter(item -> !isExcludedIp(item.getStr("ip")))
                        .collect(Collectors.toList())
        );
    }

    private boolean isExcludedIp(String ip) {
        return ip.startsWith("192.168.") || ip.startsWith("10.") || Arrays.asList(officeIps).contains(ip);
    }

    private JSONObject processItem(JSONObject item) {
        String clientIp = item.getStr("client_ip");
        String realIp = item.getByPath("request.headers.x-real-ip", String.class);
        String ip = realIp != null ? realIp : clientIp;

        List<String> tags = new ArrayList<>();
        JSONObject ipLocation = new JSONObject();

        if (isExcludedIp(ip)) {
            tags.add("local");
        } else {
            IpAddress ipAddress = ipAddressService.getIpAddress(ip);
            ipLocation = ipAddress.getIpData();

            if (!"CN".equals(ipLocation.getStr("areacode")) || StrUtil.equalsAny(ipLocation.getStr("prov"), "中国台湾", "中国香港")) {
                tags.add("overseas");
            }
        }

        if (item.getByPath("request.uri", String.class).endsWith(".php")) {
            tags.add("php");
        }

        JSONObject outputItem = new JSONObject();
        outputItem.set("host", item.getByPath("request.headers.host"));
        outputItem.set("path", item.getByPath("request.uri"));
        outputItem.set("status", item.getByPath("response.status"));
        outputItem.set("ip", ip);
        outputItem.set("ip_location", ipLocation);
        outputItem.set("tags", tags);
        outputItem.set("headers", item.getByPath("request.headers"));
        outputItem.set("time", item.getLong("start_time"));
        outputItem.set("method", item.getByPath("request.method"));
        return outputItem;
    }

    @Scheduled(cron = "0/30 * * * * ?")
    private void heartbeat() {
        broadcastMessage(new JSONObject().set("type", "heartbeat").toStringPretty());
    }
}
